# encoding = utf-8
import json
import os
import cv2 as cv
import matplotlib.pyplot as plt
import torch
from torchvision import transforms
from network.model_v3 import mobilenet_v3_large
from utils import ConfigReader

# Read Config
CR = ConfigReader('config.yaml')
num_classes = CR.getElement("num_classes")
input_size = CR.getElement("input_size")
pth_path = CR.getElement('pth_path')
img_path = CR.getElement('img_path')
# set transform
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 需要与训练出的模型相同
data_transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Resize([input_size, input_size]),
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]
)
# read class_indict
json_path = './data/dataset/runs/class_indices.json'
assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
json_file = open(json_path, "r")
class_indict = json.load(json_file)
# create network
model = mobilenet_v3_large(num_classes=num_classes, reduced_tail=True).to(device)
# load network weights
model.load_state_dict(torch.load(pth_path, map_location=device))
# read image or image dir
if os.path.isfile(img_path):
    type = 'img'
elif os.path.isdir(img_path):
    type = 'dir'


def predict():
    if type == 'dir':
        # load image
        img_list = os.listdir(img_path)
        for img_str in img_list:
            # 读取
            image = cv.imread(img_path + img_str)
            image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
            img = data_transform(image)
            # 增加输入图像的维度
            img = torch.unsqueeze(img, dim=0)
            model.eval()
            with torch.no_grad():
                # predict class
                output = torch.squeeze(model(img.to(device))).cpu()  # 各类得分
                predict = torch.softmax(output, dim=0)  # 经过softmax层得到各类概率
                predict_cla = torch.argmax(predict).numpy()

            cv.imshow("show", image)

            print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
                                                         predict[predict_cla].numpy())
            print(print_res, "   ", img_str)
            c = cv.waitKey(0) & 0xFF
            if c == 27:
                break

    elif type == 'img':
        image = cv.imread(img_path)
        image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
        img = data_transform(image)
        # 增加输入图像的维度
        img = torch.unsqueeze(img, dim=0)
        model.eval()
        with torch.no_grad():
            # predict class
            output = torch.squeeze(model(img.to(device))).cpu()  # 各类得分
            predict = torch.softmax(output, dim=0)  # 经过softmax层得到各类概率
            predict_cla = torch.argmax(predict).numpy()

            print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
                                                         predict[predict_cla].numpy())
        print(print_res)
        plt.imshow(image)
        plt.show()
        cv.imshow("show", image)

        print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
                                                     predict[predict_cla].numpy())
        print(print_res)


if __name__ == '__main__':
    predict()
